# -*- coding: utf-8 -*-
"""
Created on Thu Jul 30 17:38:10 2020

@author: 1
"""
from torch.utils.data import Dataset
from PIL import Image
import numpy as np
from readcsv import readfunc


class SeaDataset(Dataset):
    def __init__(self, data_path='./raw/', transform=None, istrain=True):
        if istrain:
            self.imgs_list, self.label_list = readfunc(data_path + 'training.csv')  # readcsv.py
        else:
            self.imgs_list, self.label_list = readfunc(data_path + 'annotation.csv')
        self.transform = transform
        self.img_path = data_path + '/data/'
        
    def __getitem__(self, index):
        image = Image.open(self.img_path + self.imgs_list[index] + '.jpg').convert('RGB')  # PIL
        label = self.label_list[index]
        label = int(label.replace('\n', ''))  # 转换为int。 为什么readfunc出来的label会有那么多换行符呢?
        if self.transform is not None:
            image = self.transform(image)
        image = np.array(image)
        image = image.astype(np.float32)  # 转换为ndarray - np.float32
        return image, label  # ndarray, int
    
    def __len__(self):
        return len(self.imgs_list)


if __name__ ==  "__main__":
    dataset = SeaDataset(istrain = True)
    print(dataset.__len__())
    shapes = [dataset.__getitem__(i)[0].shape for i in range(len(dataset))]
    max_shape = np.max(shapes, 0)  # Train:800*800  Valid:540*800
    min_shape = np.min(shapes, 0)  # Train:122*150  Valid:99*196
    print(max_shape)
    print(min_shape)
    
    image, label = dataset.__getitem__(32)
    print(image.shape)
    print(label)
